import numpy as np
import pandas as pd
import statsmodels.api as sm

# Simulated steady-state data (from previous simulation)
omegas = np.array([0.05e6, 0.1e6, 0.2e6, 0.5e6])  # Rabi frequencies (MHz)
gammas = np.array([0.05e6, 0.1e6, 0.2e6])          # Decoherence rates (MHz)
steady_state_data = []
for omega in omegas:
    for gamma in gammas:
        # Simplified steady-state calculation (placeholder, use actual data from simulation)
        steady_pop = 0.5 * (1 - (omega / (omega + gamma))**2)  # Example model
        steady_state_data.append([omega * 1e-6, gamma * 1e-6, steady_pop])

# Convert to DataFrame
df = pd.DataFrame(steady_state_data, columns=['Omega (MHz)', 'Gamma (MHz)', 'Steady_State_Excited_Population'])

# Calculate Omega/Gamma ratio
df['Omega_to_Gamma_Ratio'] = df['Omega (MHz)'] / df['Gamma (MHz)']

# Perform OLS regression
X = df['Omega_to_Gamma_Ratio']
X = sm.add_constant(X)  # Add intercept term
y = df['Steady_State_Excited_Population']
model = sm.OLS(y, X).fit()

# Extract regression results
regression_results = [
    ['Intercept', model.params[0], model.pvalues[0]],
    ['Omega_to_Gamma_Ratio', model.params[1], model.pvalues[1]]
]

# Save to CSV
regression_df = pd.DataFrame(regression_results, columns=['Variable', 'Coefficient', 'P_Value'])
regression_df.to_csv('tls_regression_results.csv', index=False)